#ifndef BOSS_Afterburner_BOSSOutputLoader_H
#define BOSS_Afterburner_BOSSOutputLoader_H

#include "Containers/BranchPlotOptions.h"
#include "File/ChainLoader.h"
#include "File/FileInventory.h"
#include "Functions/PrintTable.h"
#include "TChain.h"
#include "TString.h"
#include <list>
#include <string>
#include <unordered_map>
#include <yaml-cpp/yaml.h>

/// @addtogroup Afterburner_objects
/// @{

/// General container for BOSS analysis output
/// This object only needs a path of a directory containin ROOT files from the initial analysis or of one such root file. All `TTree`s and the branches in these files will be loaded automatically, so you don't have to worry about names of those trees. This object also comes with several methods that faciliate e.g. plotting of the `TTree` contents.
///
/// **Motivation:** To investigate the contents of the ROOT file, you first need to know which `TTree`s and branches it contains. If you simply construct the `BOSSOutputLoader` by giving it a file name, all `TTree`s will be loaded automatically as well as addresses for each of their branches. Five the constructer `true` as its second argument, and the names of these `TTree`s, their branches, and the types of these branches (behind the slash `/` after the name) will be printed to the terminal. **Do this if your macro throws an exception, because this probably means that you use the wrong names for the trees and or the branches further on in the macro.**
/// @author   Remco de Boer 雷穆克 (r.e.deboer@students.uu.nl or remco.de.boer@ihep.ac.cn)
/// @date     October 24th, 2018
class BOSSOutputLoader
{
public:
  BOSSOutputLoader() {}
  BOSSOutputLoader(const TString& directoryPath);
  BOSSOutputLoader(const YAML::Node& input);
  void Import(const TString& path);

  Bool_t HasChain(const std::string& chainName) const;
  Bool_t IsZombie() const { return !fFileNames.GetNFiles(); }

  TH1F*  Draw(const std::string& chainName, const char* branchX, const Int_t nBinx, const double x1,
              const double x2, Option_t* option = "", const TString& logScale = "",
              const char* cut = "");
  TH2F*  Draw(const std::string& chainName, const char* branchX, const char* branchY,
              const Int_t nBinx, const double x1, const double x2, const Int_t nBiny,
              const double y1, const double y2, Option_t* option = "", const TString& logScale = "",
              const char* cut = "");
  TH1*   Draw(const std::string& chainName, const char* branchNames, const char* cut = "",
              Option_t* option = "", const TString& logScale = "");
  TH1*   Draw(const BranchPlotOptions& options);
  TH1*   Draw(const YAML::Node& options);
  TH1*   Draw(const std::string& chainName, const YAML::Node& options);
  void   DrawAndSaveAllBranches(const std::string& chainName, Option_t* option = "",
                                const TString& logScale = "");
  void   QuickDrawAndSaveAll(Option_t* option = "");

  void  Print() const;
  void  PrintFiles() const;
  void  PrintChains() const;
  void  Print(const char* nameOfTree, Option_t* option = "toponly");
  void  PrintCutFlow();
  Table PrintCuts(bool print = true);
  void  PrintStructure();

  Long64_t     GetEntries(const std::string& chainName);
  Long64_t     GetLargestEntries() const;
  ChainLoader& GetChain(const std::string& chainName) { return fChains.at(chainName); }
  ChainLoader& operator[](const std::string& chainName) { return GetChain(chainName); }
  template <typename TYPE>
  inline TYPE& GetBranch(const std::string& chainName, const std::string& branchName);
  std::unordered_map<std::string, ChainLoader>& GetChains() { return fChains; }

protected:
  FileInventory                                fFileNames;
  std::unordered_map<std::string, ChainLoader> fChains;
  ///< Map of `TChain`s of the loaded ROOT files (contained withing `ChainLoader`).

  std::list<ChainLoader*> CreateOrderedMap();
  void                    LoadChains();
};

class DataSetForThreeAnalysisSteps
{
public:
  DataSetForThreeAnalysisSteps() {}
  const std::shared_ptr<BOSSOutputLoader> GetInclusiveMC() const { return fDataSet_InclusiveMC; };
  const std::shared_ptr<BOSSOutputLoader> GetExclusiveMC() const { return fDataSet_ExclusiveMC; };
  const std::shared_ptr<BOSSOutputLoader> GetMeasurements() const { return fDataSet_Measurements; };
  const std::shared_ptr<BOSSOutputLoader> ImportInclusiveMC(const TString& path)
  {
    return ImportPathToDataSet(path, fDataSet_InclusiveMC);
  }
  const std::shared_ptr<BOSSOutputLoader> ImportExclusiveMC(const TString& path)
  {
    return ImportPathToDataSet(path, fDataSet_ExclusiveMC);
  }
  const std::shared_ptr<BOSSOutputLoader> ImportMeasurements(const TString& path)
  {
    return ImportPathToDataSet(path, fDataSet_Measurements);
  }

private:
  const std::shared_ptr<BOSSOutputLoader>
    ImportPathToDataSet(const TString& path, std::shared_ptr<BOSSOutputLoader>& dataSet)
  {
    dataSet = std::make_shared<BOSSOutputLoader>(path);
    return dataSet;
  }
  std::shared_ptr<BOSSOutputLoader> fDataSet_InclusiveMC{nullptr};
  std::shared_ptr<BOSSOutputLoader> fDataSet_ExclusiveMC{nullptr};
  std::shared_ptr<BOSSOutputLoader> fDataSet_Measurements{nullptr};
};

/// @}

/// Get a certain branch with name `branchName` in TTree with name `chainName`.
/// @warning The program will `terminate` if the branch and/or tree does not exist in the file.
template <typename TYPE>
inline TYPE& BOSSOutputLoader::GetBranch(const std::string& chainName,
                                         const std::string& branchName)
{
  return GetChain(chainName).Get<TYPE>(branchName);
}

#endif